#数据处理部分之前的代码，加入部分数据处理的库
import paddle
import paddle.fluid as fluid
from paddle.fluid.dygraph.nn import FC
import numpy as np
import os
import gzip
import json
import random

# 声明数据集文件位置
datafile = '../work/mnist.json.gz'
print('loading mnist dataset from {} ......'.format(datafile))
# 加载json数据文件
data = json.load(gzip.open(datafile))
print('mnist dataset load done')
# 读取到的数据区分训练集，验证集，测试集
train_set, val_set, eval_set = data

# 数据集相关参数，图片高度IMG_ROWS, 图片宽度IMG_COLS
IMG_ROWS = 28
IMG_COLS = 28

# 打印数据信息
# imgs, label = train_set[0], train_set[1]
# print("训练数据集数量: ", len(imgs))

# # 观察验证集数量
# imgs, label = val_set[0], val_set[1]
# print("验证数据集数量: ", len(imgs))

# # 观察测试集数量
# imgs, label = val= eval_set[0], eval_set[1]
# print("测试数据集数量: ", len(imgs))

# 获得训练数据集
mode = "train"
imgs, label = train_set[0], train_set[1]
print("训练数据集数量: imgs:", len(imgs),",label:",len(label))

imgs_length = len(imgs)

assert len(imgs) == len(label), \
        "length of train_imgs({}) should be the same as train_labels({})".format(len(imgs), len(label))

# 获得训练数据集长度
imgs_length = len(imgs)
# 定义数据集每个数据的序号，根据序号读取数据
index_list = list(range(imgs_length))
# 读入数据时用到的批次大小
BATCHSIZE = 100

# 定义数据生成器
def data_generator():
    if mode == 'train':
        # 训练模式下打乱数据
        random.shuffle(index_list)
    imgs_list = []
    labels_list = []
    for i in index_list:
        # 将数据处理成希望的格式，比如类型为float32，shape为[1, 28, 28]
        img = np.reshape(imgs[i], [1, IMG_ROWS, IMG_COLS]).astype('float32')
        label = np.reshape(labels[i], [1]).astype('float32')
        imgs_list.append(img) 
        labels_list.append(label)
        if len(imgs_list) == BATCHSIZE:
            # 获得一个batchsize的数据，并返回
            yield np.array(imgs_list), np.array(labels_list)
            # 清空数据读取列表
            imgs_list = []
            labels_list = []

    # 如果剩余数据的数目小于BATCHSIZE，
    # 则剩余数据一起构成一个大小为len(imgs_list)的mini-batch
    if len(imgs_list) > 0:
        yield np.array(imgs_list), np.array(labels_list)

